import json
import time
from pprint import pprint

from src.find_kv_and_check_word.datamodel import CheckModel
from src.find_kv_and_check_word.prompt import Prompt
from src.find_kv_and_check_word.tools_hyc_cyy import clean_json_string, model_multi_2nd_output_expand_keys
from src.find_kv_and_check_word.tools_zwy_lfw import append_data
from src.find_kv_and_check_word.word_operation_generation import get_textDto, get_textDto_singe_file

from src.huey_init._huey import huey
from src.tools.chain import gen_chain
from src.tools.dms_operator import dmsoperator
from src.tools.information_exchange import pub_gen


@huey.task(context=True)
def run_check_task(check_model: CheckModel, task=None):
    """
    根据一致性检查类型，选择单文档检查或多文档互检
    """
    consistency_type = check_model.consistency_type
    if consistency_type == 1:
        for i in check_model.consistency_check_file_list:
            try:
                #保存、复制原文档信息
                original_file_info = i["consistency_check_file_info"].copy()
                check_single_file(check_model, i["consistency_check_file_info"], task)
            except Exception as e:
                # 捕获可能的错误
                print(f"An error occurred: {e}")
                # 返回未编辑的单文档
                exchange_info = build_exchange_info(check_model, original_file_info, original_file_info["realPath"], task.id)
                dmsoperator.post_exchange_info(exchange_info)
                print(f"发生错误，返回未编辑的单文档信息。")
            finally:
                # 无论是否出错都会执行的代码块
                print(f"“单文档一致性检查”运行完毕！")
    else:
        try:
            # 保存、复制原文档信息
            original_main_file_info = check_model.consistency_check_file_list[0]["consistency_check_file_info"].copy()
            check_multi_file(check_model, task)
        except Exception as e:
            # 捕获可能的错误
            print(f"An error occurred: {e}")
            # 返回未编辑的主文档
            exchange_info = build_exchange_info(check_model, original_main_file_info, original_main_file_info["realPath"], task.id)
            dmsoperator.post_exchange_info(exchange_info)
            print(f"发生错误，返回未编辑的主文档信息。")
        finally:
            # 无论是否出错都会执行的代码块
            print(f"“多文档一致性检查”运行完毕！")

def check_single_file(check_model: CheckModel, file_info, task):
    """
    处理单文档的一致性检查
    """
    start_time = time.time()
    print("单文档自检程序开始运行！")
    single_file_flag = True

    # 获取文件信息
    # file_info = check_model.consistency_check_file_list[0]["consistency_check_file_info"]
    file_realPath = file_info["realPath"]
    keyword_list = file_info["indicatorNameList"]
    
    # 生成数据库关键词列表
    keywords_from_database = ",".join(keyword_list)

    # 构造检查数据(需要将 "file_realPath" 构造成列表)
    file_realPath_list = [file_realPath]
    final = append_data(file_realPath_list, keywords_from_database, single_file_flag, task.id, file_info)
    result = get_textDto_singe_file(file_realPath, final, task.id, file_info)

    # 提交操作文件结果
    res = dmsoperator.post_operate_file(result)
    
    # 构造文档信息
    doc_info = {
        "docId": file_info["docId"],
        "taskDocId": file_info["taskDocId"],
        "documentName": file_info["documentName"],
        "realPath": res["data"],
        "indicatorNameList": file_info["indicatorNameList"]
    }

    # 交换信息
    exchange_info = build_exchange_info(check_model, file_info, res["data"], task.id)
    print("res-data----:",res["data"])
    dmsoperator.post_exchange_info(exchange_info)

    # 输出执行时间
    end_time = time.time()
    print_execution_time(start_time, end_time)


def check_multi_file(check_model: CheckModel, task):
    """
    处理多文档的一致性互检
    """
    start_time = time.time()
    print("多文档互检程序开始运行！")
    single_file_flag = False

    # 获取主文件信息
    main_file_info = check_model.consistency_check_file_list[0]["consistency_check_file_info"]
    main_file_realPath = main_file_info["realPath"]
    keyword_list = main_file_info["indicatorNameList"]
    
    # 获取参考文件信息
    file_id_list = [main_file_realPath]
    refer_file_info = check_model.consistency_refer_file_list
    document_name = {item['realPath']: item['documentName'] for item in refer_file_info}

    # 构建所有文件的ID列表
    file_id_list.extend([refer_file["realPath"] for refer_file in refer_file_info])

    # 生成数据库关键词列表
    keywords_from_database = ",".join(keyword_list)

    # 构造检查数据
    final = append_data(file_id_list, keywords_from_database, single_file_flag, task.id, main_file_info)
    
    # 如果多文档互检时，主文档未找到关键字，直接返回未编辑的文档
    if final and isinstance(final[0], list) and not final[0]:
        # 返回未编辑的主文档
        exchange_info = build_exchange_info(check_model, main_file_info, main_file_info["realPath"], task.id)
        dmsoperator.post_exchange_info(exchange_info)
    else:
        # 获取相似关键词提示
        prompt = Prompt()
        data_test = final[0]
        context = prompt.find_same_mean_keyword_prompt_multiple_files(data_test, main_file_realPath)
        result = json.loads(clean_json_string(gen_chain(context)))
        # （大模型输出速度优化）键名扩展
        result = model_multi_2nd_output_expand_keys(result)
        # 如果大模型返回一个空列表（即多文件互检没有对比结果），则返回未编辑的主文档
        if result['result'] == []:
            exchange_info = build_exchange_info(check_model, main_file_info, main_file_info["realPath"], task.id)
            dmsoperator.post_exchange_info(exchange_info)

        # 有对比结果
        else: 
            # 发布生成的结果
            pub_gen(task.id, str(result), main_file_info.get("docId"), main_file_info.get("dmsDocId"))
            pprint(result)

            # 提交操作文件结果
            res = dmsoperator.post_operate_file(get_textDto(main_file_realPath, result, final[1], document_name))
            pprint(res)

            # 交换信息
            exchange_info = build_exchange_info(check_model, main_file_info, res["data"], task.id)
            dmsoperator.post_exchange_info(exchange_info)

        # 输出执行时间
        end_time = time.time()
        print_execution_time(start_time, end_time)


def build_exchange_info(check_model, file_info, result_data, task_id):
    """
    构建交换信息的字典
    """
    return {
        "functionType": check_model.exchange_info.get("functionType"),
        "taskType": check_model.exchange_info.get("taskType"),
        "taskDocumentType": check_model.exchange_info.get("taskDocumentType"),
        "executionType": check_model.exchange_info.get("executionType"),
        "isSuccess": True,
        "dmsTaskId": check_model.exchange_info.get("dmsTaskId"),
        "docId": file_info["docId"],
        "taskDocId": file_info["taskDocId"],
        "resultFirstFilePath": result_data,
        "aiTaskId": task_id
    }


def print_execution_time(start_time, end_time):
    """
    打印任务执行时间
    """
    execution_time = end_time - start_time
    print(f"执行时间: {execution_time:.6f} 秒")